# coding:utf-8
import numpy as np
import pandas as pd
from collections import *
import csv, os, sys
from sklearn.ensemble import RandomForestClassifier
from sklearn.externals import joblib

# N-gram 
def getOpcodeNgram(ops , n):
    opngramlist = [''.join(ops[i:i + n]) for i in range(len(ops) - n)]
    opngram = Counter(opngramlist)
    return opngram

# 获取特征
def getFeature():
    dirPath = 'data'
    totalFea = []
    for cg in os.listdir(dirPath):
        path = dirPath + '/' + cg
        totalCounts = Counter()
        for i in range(10):
            p = path + '/' + str(i) + '.txt'
            with open(p, 'r') as f:
                content = f.read().replace(' ', '').replace('\n', '').replace('\t', '')
                counts = getOpcodeNgram(content, 3)
                totalCounts.update(counts)
        feas = totalCounts.most_common(50)
        totalFea.extend(feas)
    features = [e[0] for e in totalFea]
    with open('result/feature.csv', 'w') as f:
        writer = csv.writer(f)
        writer.writerow(features)

# 获取数据               
def getData():
    with open('result/feature.csv', 'r') as f:
        feas = list(csv.reader(f))[0]
    dirPath = 'data'
    for cg in os.listdir(dirPath):
        path = dirPath + '/' + cg
        for i in range(10):
            p = path + '/' + str(i) + '.txt'
            print(p)
            data = [0 for e in feas]
            for c in feas:
                with open(p, 'r') as ff:
                    content = ff.read()
                    ops = dict(getOpcodeNgram(content, 3))
                    length = len(content)
                    if c in ops.keys():
                        data[feas.index(c)] += ops[c]
            data = [e / length for e in data]
            data.extend([cg])
            with open('result/data.csv', 'a+', newline='') as f:
                writer = csv.writer(f)
                writer.writerow(data)
                
# 获取模型                
def getModel():
    with open('result/data.csv', 'r') as f:
        content = list(csv.reader(f))
    trainData = [e[:-1] for e in content]
    trainLabel = [e[-1] for e in content]
    clf = RandomForestClassifier(n_estimators=100).fit(trainData, trainLabel)
    joblib.dump(clf, "result/train_model.m")

# 通过文件内容预测
def PredictByFile(path):
    with open('result/feature.csv', 'r') as f:
        feas = list(csv.reader(f))[0]
    data = [0 for e in feas]
    with open(path, 'r')as f:
        content = f.read()       
        ops = dict(getOpcodeNgram(content, 3))
        length = len(content)
        for c in feas:
            if c in ops.keys():
                data[feas.index(c)] += ops[c]
        if length == 0:
            print("Error! File is empty!")
        data = [e / length for e in data]
    clf = joblib.load('result/train_model.m')
    PredRes = clf.predict([data])
    print(PredRes)

# 通过字符串内容预测    
def PredictByString(content):
    with open('python/feature.csv', 'r') as f:
        feas = list(csv.reader(f))[0] 
    data = [0 for e in feas]     
    ops = dict(getOpcodeNgram(content, 3))
    length = len(content)
    for c in feas:
        if c in ops.keys():
            data[feas.index(c)] += ops[c]
    if length == 0:
        print("Error! File is empty!")
    data = [e / length for e in data]
    clf = joblib.load('python/train_model.m')
    PredRes = clf.predict([data])
    return PredRes[0]


content = sys.argv[1]
print(PredictByString(content))
     
     
     
     
